#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <memory>

#include "MyMD5.h"
#include "AES.h"


/*
size_t fread(void *ptr, size_t size_of_elements,
             size_t number_of_elements, FILE *a_file);

size_t fwrite(const void *ptr, size_t size_of_elements,
             size_t number_of_elements, FILE *a_file);
*/



void file_encrypt(const char* infile, const char* outfile, const uint8_t *key) {
    FILE* fp = fopen(infile,  "rb");
    FILE* op = fopen(outfile, "wb");

    uint8_t ct[16] = { 0 };     // [out] 外部申请输出数据内存，用于加密后的数据
    uint8_t plain[16] = { 0 };  // [out] 外部申请输出数据内存，用于解密后的数据

    uint8_t buf[16];
    uint8_t PAD = 0;

    int ret = fread(buf, 1, 16, fp);
    printf("ret=%d\n", ret);
    while (ret > 0) {
        if (ret != 16) {
            for (int i = ret; i < 16; i++) {
                buf[i] = PAD;
            }
        }
        aesEncrypt(key, 16, buf, ct, 16);  // 加密 ct是结果

        fwrite(ct, 1, 16, op);

        ret = fread(buf, 1, 16, fp); 
    } 
     
    fclose(fp);
    fclose(op);
    printf("encry succ in %s\n", outfile);
}


void file_decrypt(const char* infile, const char* outfile, const uint8_t* key) {
    FILE* fp = fopen(infile, "rb");
    FILE* op = fopen(outfile, "wb");

    uint8_t ct[16] = { 0 };     // [out] 外部申请输出数据内存，用于加密后的数据
    uint8_t plain[16] = { 0 };  // [out] 外部申请输出数据内存，用于解密后的数据

    uint8_t buf[16];
    uint8_t preBuf[16];
    uint8_t PAD = 0;

    int ret = fread(buf, 1, 16, fp); 
    while (ret > 0) {
        if (ret != 16) {
            printf("你的文件有错误，不是16位对齐的\n!");
            break;
        } 
        memmove(preBuf, buf, 16);
        ret = fread(buf, 1, 16, fp);

        if (ret < 1) { //最后的一个了
            aesDecrypt(key, 16, preBuf, plain, 16);       // 解密  plain是结果 
             

            // 移除多余的0
            int i = 15;
            while(i>=0 && plain[i]==0) {
                  i--;
            }

            fwrite(plain, 1, i+1, op);
            break;
        }
        else {
            aesDecrypt(key, 16, preBuf, plain, 16);       // 解密  plain是结果 
            fwrite(plain, 1, 16, op);
        } 

    } 
    fclose(fp);
    fclose(op);
    printf("encry succ in %s\n", outfile);
}


void test_md5() {
    unsigned char pwd[256];
    unsigned char decrypt[16];
    int i;
    printf(">>");
    scanf("%s", pwd);
    MD5_CTX md5;
    MD5Init(&md5);
    MD5Update(&md5, pwd, strlen((char*)pwd));
    MD5Final(&md5, decrypt);
    printf("加密前:%s\n加密后32位:", pwd);
    for (i = 0; i < 16; i++)
    {
        printf("%02x", decrypt[i]);
    }
    printf("\n加密后16位：");
    for (int i = 4; i < 12; i++) {
        printf("%02x", decrypt[i]);
    } 
}


void get_md5(uint8_t* res) {
    unsigned char pwd[256];
    unsigned char decrypt[16];
    int i;

    scanf("%s", pwd);
    MD5_CTX md5;
    MD5Init(&md5);
    MD5Update(&md5, pwd, strlen((char*)pwd));
    MD5Final(&md5, decrypt); 
    

    memmove(res, decrypt, 16);

}

bool check_file_exist(const char* fn) {
    FILE* fp =  fopen(fn, "rb");
    bool ret=true;
    if (!fp){
        ret= false;
    } 
    fclose(fp);
    return ret;
}


int main() {
    // const uint8_t key[16] = { 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c };
    uint8_t key[16] = {0};
    bool hasInputPwd = false;

    const char* ori_file = "data.txt";
    const char* enc_file = "data.lock";


    printf("【这是一个基于AES加密的密码管理器】\n");
    printf("文件[%s] 存在=[%d]\n", ori_file, check_file_exist(ori_file));
    printf("文件[%s] 存在=[%d]\n", enc_file, check_file_exist(enc_file));
    int operate = 0;
    while (true) {
        printf("请输入对应的操作：1读取，2保存，3修改密码, 0退出:");
        scanf("%d", &operate);
        if (operate == 1) {
            if ( ! check_file_exist(enc_file)) {
                printf("文件[%s]不存在！无法读取\n", enc_file);
                continue;
            } 
            if (!hasInputPwd) {
                printf("输入解密密码：");
                get_md5(key);
                hasInputPwd = true;
            }  
            file_decrypt(enc_file, ori_file, key);
            char cmd[256];
            sprintf_s(cmd, "start %s", ori_file);
            system(cmd);
        }
        else if (operate == 2) {
             if ( ! check_file_exist(ori_file)) {
                printf("文件[%s]不存在！无法保存\n", ori_file);
                continue;
            } 
            if (!hasInputPwd) {
                printf("输入加密密码：");
                get_md5(key);
                hasInputPwd = true;
            } 
            file_encrypt(ori_file, enc_file, key);
            printf("remove ret=%d\n", remove(ori_file));
            printf("保存成功\n");
        }
        else if (operate == 0) {
            break;
        }
        else if (operate == 3) {
            printf("输入密码：");
            get_md5(key);
            hasInputPwd = true;
            printf("修改成功\n");
        }
        else {
            printf("操作错误\n");
        }
    }
    
     

    return 0;
}